package com.xjgz.aio.socket.pro.extension.ssl.factory;

import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.*;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Date;

/**
 * 自签名SSL证书工厂类，用于自动生成服务端SSL证书
 */
public class AutoServerSSLContextFactory implements SSLContextFactory {
    private final String commonName;

    private final String organization;

    private final String organizationUnit;

    private final int validityDays;

    private final String keyAlgorithm;

    private final int keySize;

    /**
     * @param commonName       通用名称（CN）
     * @param organization     组织名称（O）
     * @param organizationUnit 组织单元名称（OU）
     * @param validityDays     证书有效期（天）
     * @param keyAlgorithm     密钥算法（RSA、DSA、EC等）
     * @param keySize          密钥长度
     */
    public AutoServerSSLContextFactory(
            String commonName,
            String organization,
            String organizationUnit,
            int validityDays,
            String keyAlgorithm,
            int keySize
    ) {
        this.commonName = commonName;
        this.organization = organization;
        this.organizationUnit = organizationUnit;
        this.validityDays = validityDays;
        this.keyAlgorithm = keyAlgorithm;
        this.keySize = keySize;
    }

    public AutoServerSSLContextFactory() {
        this(
                "localhost",
                "learning",
                "java learning",
                365,
                "RSA",
                2048
        );
    }

    @Override
    public SSLContext create() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(keyAlgorithm);
        keyPairGenerator.initialize(keySize, new SecureRandom());
        KeyPair keyPair = keyPairGenerator.generateKeyPair();

        PrivateKey privateKey = keyPair.getPrivate();
        // 生成自签名证书
        X509Certificate certificate = generateSelfSignedCertificate(keyPair);

        KeyStore keyStore = KeyStore.getInstance("JKS");
        keyStore.load(null, null);
        Certificate[] chain = {certificate};
        keyStore.setKeyEntry("server", privateKey, new char[0], chain);

        KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance("SunX509");
        keyManagerFactory.init(keyStore, null);

        SSLContext sslContext = SSLContext.getInstance("TLS");
        sslContext.init(keyManagerFactory.getKeyManagers(), null, new SecureRandom());
        return sslContext;
    }

    @Override
    public void initSSLEngine(SSLEngine sslEngine) throws Exception {
        sslEngine.setUseClientMode(false);
    }

    private X509Certificate generateSelfSignedCertificate(KeyPair keyPair) throws Exception {
        // 设置主题和颁发者
        String dn = "CN=" + commonName + ", O=" + organization + ", OU=" + organizationUnit;
        X500Name issuer = new X500Name(dn);
        X500Name subject = new X500Name(dn);

        // 设置序列号
        BigInteger serialNumber = new BigInteger(64, new SecureRandom());

        // 设置有效期
        long currentTimeMillis = System.currentTimeMillis();
        Date notBefore = new Date(currentTimeMillis);
        Date notAfter = new Date(currentTimeMillis + validityDays * 24L * 60 * 60 * 1000);

        // 创建证书构建器
        X509v3CertificateBuilder certBuilder = new JcaX509v3CertificateBuilder(
                issuer,
                serialNumber,
                notBefore,
                notAfter,
                subject,
                keyPair.getPublic()
        );

        // 添加基本约束扩展（CA证书）
        certBuilder.addExtension(
                Extension.basicConstraints,
                true,
                new BasicConstraints(true)
        );

        // 添加密钥用法扩展
        certBuilder.addExtension(
                Extension.keyUsage,
                true,
                new KeyUsage(KeyUsage.digitalSignature | KeyUsage.keyEncipherment | KeyUsage.keyCertSign)
        );

        // 添加主题替代名称扩展
        GeneralName[] generalNames = new GeneralName[]{
                new GeneralName(GeneralName.dNSName, commonName)
        };
        certBuilder.addExtension(
                Extension.subjectAlternativeName,
                false,
                new GeneralNames(generalNames)
        );

        // 签名证书
        ContentSigner signer = new JcaContentSignerBuilder("SHA256withRSA")
                .build(keyPair.getPrivate());

        // 生成证书
        return new JcaX509CertificateConverter()
                .getCertificate(certBuilder.build(signer));
    }

}
